package com.ethink.gateway.filter.pre;

import com.ethink.framework.common.util.JacksonUtils;
import com.ethink.framework.common.util.PathMatcherUtil;
import com.ethink.framework.log.constant.LogConstant;
import com.ethink.gateway.constant.GatewayConstant;
import com.ethink.gateway.filter.AbstractPreFilter;
import com.ethink.gateway.model.RequestInfo;
import com.ethink.gateway.properties.GateWayCommonProperties;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.collections4.MapUtils;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.MDC;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
import org.springframework.http.MediaType;
import org.springframework.http.codec.HttpMessageReader;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.stereotype.Component;
import org.springframework.web.reactive.function.server.HandlerStrategies;
import org.springframework.web.reactive.function.server.ServerRequest;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.net.InetSocketAddress;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

@Slf4j
@Component
public class RequestLogFilter extends AbstractPreFilter {
    private static final String CLIENT_IP_HEADER_KEY = "X-Forwarded-For";
    private static final String CLIENT_PORT_HEADER_KEY = "X-Forwarded-Port";
    @Autowired
    private GateWayCommonProperties commonProperties;

    private static final List<HttpMessageReader<?>> messageReaders =
            HandlerStrategies.withDefaults().messageReaders();

    @Override
    public int getOrder() {
        return GatewayConstant.FILTER_ORDER_REQUEST_LOG;
    }

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        ServerHttpRequest request = exchange.getRequest();
        RequestInfo requestInfo = buildRequestLog(request);
        exchange.getAttributes().put(GatewayConstant.ATTR_REQUEST_INFO, requestInfo);
        if (PathMatcherUtil.notMatch(commonProperties.getRequestLogExcludePaths(), requestInfo.getUri())) {
            log.info("request: {}", JacksonUtils.toJson(requestInfo));
            if (request.getMethod() != HttpMethod.GET) {
                MediaType contentType = request.getHeaders().getContentType();
                if (MediaType.APPLICATION_JSON.equals(contentType) || MediaType.APPLICATION_JSON_UTF8.equals(contentType)) {
                    return readBody(exchange, chain);
                }
            }
        }
        return chain.filter(exchange);
    }


    @Override
    protected ServerWebExchange preFilter(ServerWebExchange exchange) {
        return exchange;
    }

    /**
     * ReadJsonBody
     *
     * @param exchange
     * @param chain
     * @return
     */
    private Mono<Void> readBody(ServerWebExchange exchange, GatewayFilterChain chain) {
        /**
         * join the body
         */
        return DataBufferUtils.join(exchange.getRequest().getBody())
                .flatMap(dataBuffer -> {
                    byte[] bytes = new byte[dataBuffer.readableByteCount()];
                    dataBuffer.read(bytes);
                    DataBufferUtils.release(dataBuffer);
                    Flux<DataBuffer> cachedFlux = Flux.defer(() -> {
                        DataBuffer buffer =
                                exchange.getResponse().bufferFactory().wrap(bytes);
                        DataBufferUtils.retain(buffer);
                        return Mono.just(buffer);
                    });
                    ServerHttpRequest mutatedRequest =
                            new ServerHttpRequestDecorator(exchange.getRequest()) {
                                @Override
                                public Flux<DataBuffer> getBody() {
                                    return cachedFlux;
                                }
                            };
                    ServerWebExchange mutatedExchange =
                            exchange.mutate().request(mutatedRequest).build();
                    return ServerRequest.create(mutatedExchange, messageReaders)
                            .bodyToMono(String.class)
                            .doOnNext(objectValue -> {
                                log.info("requestBody: {}", objectValue);
                            }).then(chain.filter(mutatedExchange));
                });
    }

    private RequestInfo buildRequestLog(ServerHttpRequest request) {
        String traceId = MDC.get(LogConstant.LOG_TRACE_ID);
        InetSocketAddress address = request.getRemoteAddress();
        HttpMethod method = request.getMethod();
        HttpHeaders headers = request.getHeaders();
        Map<String, String> requestHeaders = headers.toSingleValueMap();
        Map<String, String> headerMap = new HashMap<>();
        headerMap.putAll(requestHeaders);
        if (MapUtils.isNotEmpty(requestHeaders) && CollectionUtils.isNotEmpty(commonProperties.getNotLogHeaders())) {
            for (String header : commonProperties.getNotLogHeaders()) {
                headerMap.remove(header);
            }
        }
        String portStr = request.getHeaders().getFirst(CLIENT_PORT_HEADER_KEY);
        int clientPort = StringUtils.isNumeric(portStr) ? Integer.valueOf(portStr) : -1;
        RequestInfo requestInfo = new RequestInfo();
        requestInfo.setClientIp(request.getHeaders().getFirst(CLIENT_IP_HEADER_KEY));
        requestInfo.setClientPort(clientPort);
        requestInfo.setRequestIp(address.getHostName());
        requestInfo.setRequestPort(address.getPort());
        requestInfo.setUrl(request.getURI().toString());
        requestInfo.setUri(request.getPath().value());
        requestInfo.setMethod(method.name());
        requestInfo.setQueryString(request.getQueryParams().toSingleValueMap().toString());
        requestInfo.setHeaders(headerMap);
        requestInfo.setTraceId(traceId);
        requestInfo.setRequestTime(System.currentTimeMillis());
        return requestInfo;
    }
}
